#include <optional>
#include <sstream>
#include <string>

#include "common/config.h"
#include "common/exception.h"
#include "common/logger.h"
#include "common/rid.h"
#include "storage/index/b_plus_tree.h"
#include "storage/page/b_plus_tree_header_page.h"
#include "storage/page/b_plus_tree_page.h"
#include "storage/page/page_guard.h"

namespace bustub {

INDEX_TEMPLATE_ARGUMENTS
BPLUSTREE_TYPE::BPlusTree(std::string name, page_id_t header_page_id, BufferPoolManager *buffer_pool_manager,
                          const KeyComparator &comparator, int leaf_max_size, int internal_max_size)
    : index_name_(std::move(name)),
      bpm_(buffer_pool_manager),
      comparator_(std::move(comparator)),
      leaf_max_size_(leaf_max_size),
      internal_max_size_(internal_max_size),
      header_page_id_(header_page_id) {
  WritePageGuard guard = bpm_->FetchPageWrite(header_page_id_);
  auto root_page = guard.AsMut<BPlusTreeHeaderPage>();
  root_page->root_page_id_ = INVALID_PAGE_ID;
}

/*
 * Helper function to decide whether current b+tree is empty
 */
INDEX_TEMPLATE_ARGUMENTS
auto BPLUSTREE_TYPE::IsEmpty() const -> bool {
  ReadPageGuard guard = bpm_->FetchPageRead(header_page_id_);
  auto header_page = guard.As<BPlusTreeHeaderPage>();
  return header_page->root_page_id_ == INVALID_PAGE_ID;
}
/*****************************************************************************
 * SEARCH
 *****************************************************************************/
/*
 * Return the only value that associated with input key
 * This method is used for point query
 * @return : true means key exists
 */
INDEX_TEMPLATE_ARGUMENTS
auto BPLUSTREE_TYPE::GetValue(const KeyType &key, std::vector<ValueType> *result, Transaction *txn) -> bool {
  if (IsEmpty()) {
    return false;
  }
  std::vector<ValueType> ans;
  Context ctx;
  if (not FindLeafRead(ctx, key)) {
    return false;
  }

  bool flag = false;
  auto base_page = ctx.read_set_.back().As<BPlusTreePage>();
  if (base_page->IsLeafPage()) {
    int idx = FindIdx(base_page, key);
    auto page = ctx.read_set_.back().As<LeafPage>();
    if (idx > 0 && comparator_(page->KeyAt(idx - 1), key) == 0) {
      ans.push_back(page->ValueAt(idx - 1));
      flag = true;
    }
  }

  if (flag) {
    *result = std::move(ans);
  }
  CleanUp(ctx);

  return flag;
}

/*****************************************************************************
 * INSERTION
 *****************************************************************************/
/*
 * Insert constant key & value pair into b+ tree
 * if current tree is empty, start new tree, update root page id and insert
 * entry, otherwise insert into leaf page.
 * @return: since we only support unique key, if user try to insert duplicate
 * keys return false, otherwise return true.
//  */
INDEX_TEMPLATE_ARGUMENTS
auto BPLUSTREE_TYPE::Insert(const KeyType &key, const ValueType &value, Transaction *txn) -> bool {
  if (IsEmpty()) {
    WritePageGuard header_guard = bpm_->FetchPageWrite(header_page_id_);
    page_id_t root_id = header_guard.As<BPlusTreeHeaderPage>()->root_page_id_;
    if (root_id == INVALID_PAGE_ID) {
      BasicPageGuard guard = bpm_->NewPageGuarded(&root_id);
      auto page = guard.AsMut<LeafPage>();
      page->SetNextPageId(INVALID_PAGE_ID);
      page->Init(leaf_max_size_);
      page->InsertAt(0, MappingType(key, value));

      header_guard.AsMut<BPlusTreeHeaderPage>()->root_page_id_ = root_id;
      return true;
    }
  }
  // optimistic find
  Context ctx;
  if (not OptimisticFindLeaf(ctx, key)) {
    return false;
  }
  auto page = ctx.write_set_.back().AsMut<LeafPage>();
  int idx = FindIdx(page, key);
  if (idx > 0 && comparator_(page->KeyAt(idx - 1), key) == 0) {
    return false;
  } else if (page->GetSize() + 1 < page->GetMaxSize()) {
    page->InsertAt(idx, MappingType(key, value));
    return true;
  }

  // restart over
  CleanUp(ctx);
  if (not FindLeafWrite(ctx, key, true)) {
    return false;
  }
  bool succeed = true;
  auto leaf_page = ctx.write_set_.back().As<LeafPage>();
  int id = FindIdx(leaf_page, key);
  if (id > 0 && comparator_(leaf_page->KeyAt(id - 1), key) == 0) {
    succeed = false;
  } else {
    InsertAtLeaf(ctx, id, key, value);
  }
  CleanUp(ctx);
  return succeed;
}

/*****************************************************************************
 * REMOVE
 *****************************************************************************/
/*
 * Delete key & value pair associated with input key
 * If current tree is empty, return immediately.
 * If not, User needs to first find the right leaf page as deletion target, then
 * delete entry from leaf page. Remember to deal with redistribute or merge if
 * necessary.
 */
INDEX_TEMPLATE_ARGUMENTS
void BPLUSTREE_TYPE::Remove(const KeyType &key, Transaction *txn) {
  if (IsEmpty()) {
    return;
  }
  // optimistic find
  Context ctx;
  if (not OptimisticFindLeaf(ctx, key)) {
    return;
  }
  auto page = ctx.write_set_.back().AsMut<LeafPage>();
  int idx = FindIdx(page, key);
  if (idx <= 0 || comparator_(page->KeyAt(idx - 1), key) != 0) {
    return;
  }
  if (page->GetSize() > page->GetMinSize()) {
    page->Remove(idx - 1);
    return;
  }

  // restart over
  CleanUp(ctx);
  if (not FindLeafWrite(ctx, key, false)) {
    return;
  }
  DeleteEntryLeaf(ctx, key);
  CleanUp(ctx);
}

INDEX_TEMPLATE_ARGUMENTS
void BPLUSTREE_TYPE::CleanUp(Context &ctx) {
  ctx.header_page_ = std::nullopt;
  while (not ctx.read_set_.empty()) {
    ctx.read_set_.front().Drop();
    ctx.read_set_.pop_front();
  }
  while (not ctx.write_set_.empty()) {
    ctx.write_set_.front().Drop();
    ctx.write_set_.pop_front();
  }
}

INDEX_TEMPLATE_ARGUMENTS
void BPLUSTREE_TYPE::DeleteEntryInternal(Context &ctx, const KeyType &key) {
  auto guard = std::move(ctx.write_set_.back());
  ctx.write_set_.pop_back();

  auto base_page = guard.As<BPlusTreePage>();
  int idx = FindIdx(base_page, key);
  auto page = guard.AsMut<InternalPage>();
  if (idx > 0 && comparator_(page->KeyAt(idx - 1), key) == 0) {
    --idx;
  } else {
    return;
  }
  page->Remove(idx);

  int min_size = page->GetMinSize();
  int page_id = guard.PageId();

  if (ctx.write_set_.empty() and page->GetSize() == 1) {
    page_id_t new_root_id = page->ValueAt(0);
    ctx.header_page_->AsMut<BPlusTreeHeaderPage>()->root_page_id_ = new_root_id;
    guard.Drop();
    bpm_->DeletePage(ctx.root_page_id_);
    ctx.root_page_id_ = new_root_id;

  } else if (page->GetSize() < min_size) {
    auto parent = ctx.write_set_.back().As<InternalPage>();
    int cur_ind = FindIdx(parent, key) - 1;
    int sib_ind = cur_ind + 1;
    KeyType mid_key = parent->KeyAt(sib_ind);
    if (cur_ind == parent->GetSize() - 1) {
      sib_ind = cur_ind - 1;
      mid_key = parent->KeyAt(cur_ind);
    }
    int sib_page_id = parent->ValueAt(sib_ind);
    auto sib_guard = bpm_->FetchPageWrite(sib_page_id);
    auto sib_page = sib_guard.AsMut<InternalPage>();

    if (sib_page->GetSize() + page->GetSize() <= internal_max_size_) {
      if (cur_ind < sib_ind) {
        sib_page->Front().first = mid_key;
        std::copy(sib_page->Begin(), sib_page->End(), page->End());
        page->IncreaseSize(sib_page->GetSize());
        DeleteEntryInternal(ctx, mid_key);
        sib_guard.Drop();
        bpm_->DeletePage(sib_page_id);
      } else {
        page->Front().first = mid_key;
        std::copy(page->Begin(), page->End(), sib_page->End());
        sib_page->IncreaseSize(page->GetSize());
        DeleteEntryInternal(ctx, mid_key);
        guard.Drop();
        bpm_->DeletePage(page_id);
      }
    } else {
      // early release
      ctx.header_page_ = std::nullopt;
      while (ctx.write_set_.size() > 1) {
        ctx.write_set_.front().Drop();
        ctx.write_set_.pop_front();
      }
      if (cur_ind < sib_ind) {
        auto &first_kv = sib_page->Front();
        sib_page->KeyAt(0) = mid_key;
        page->InsertAt(page->GetSize(), first_kv);
        sib_page->Remove(0);
        parent->KeyAt(sib_page_id) = sib_page->KeyAt(0);
      } else {
        auto &last_kv = sib_page->Back();
        page->KeyAt(0) = mid_key;
        page->InsertAt(0, last_kv);
        sib_page->Remove(sib_page->GetSize());
        parent->KeyAt(page_id) = page->KeyAt(0);
      }
    }
  }
}

INDEX_TEMPLATE_ARGUMENTS
void BPLUSTREE_TYPE::DeleteEntryLeaf(Context &ctx, const KeyType &key) {
  if (ctx.write_set_.empty()) {
    return;
  }
  auto guard = std::move(ctx.write_set_.back());
  ctx.write_set_.pop_back();

  auto base_page = guard.As<BPlusTreePage>();
  int idx = FindIdx(base_page, key);
  auto page = guard.AsMut<LeafPage>();
  if (idx > 0 && comparator_(page->KeyAt(idx - 1), key) == 0) {
    --idx;
  } else {
    return;
  }
  page->Remove(idx);

  int min_size = page->GetMinSize();
  int page_id = guard.PageId();
  if (page_id == ctx.root_page_id_) {
    if (page->GetSize() == 0) {
      ctx.header_page_->AsMut<BPlusTreeHeaderPage>()->root_page_id_ = INVALID_PAGE_ID;
      ctx.root_page_id_ = INVALID_PAGE_ID;
    }
  } else if (page->GetSize() < min_size) {
    auto parent = ctx.write_set_.back().As<InternalPage>();
    int cur_ind = FindIdx(parent, key) - 1;
    int sib_ind = cur_ind + 1;
    if (cur_ind == parent->GetSize() - 1) {
      sib_ind = cur_ind - 1;
    }
    int sib_page_id = parent->ValueAt(sib_ind);
    auto sib_guard = bpm_->FetchPageWrite(sib_page_id);
    auto sib_page = sib_guard.AsMut<LeafPage>();

    // merge to the left node
    if (sib_page->GetSize() + page->GetSize() <= leaf_max_size_) {
      if (cur_ind < sib_ind) {
        std::copy(sib_page->Begin(), sib_page->End(), page->End());
        page->SetNextPageId(sib_page->GetNextPageId());
        DeleteEntryInternal(ctx, parent->KeyAt(sib_ind));
        sib_guard.Drop();
        bpm_->DeletePage(sib_page_id);
      } else {
        std::copy(page->Begin(), page->End(), sib_page->End());
        sib_page->SetNextPageId(page->GetNextPageId());
        DeleteEntryInternal(ctx, parent->KeyAt(cur_ind));
        guard.Drop();
        bpm_->DeletePage(page_id);
      }
    } else {
      // early release
      ctx.header_page_ = std::nullopt;
      while (ctx.write_set_.size() > 1) {
        ctx.write_set_.front().Drop();
        ctx.write_set_.pop_front();
      }
      if (cur_ind < sib_ind) {
        MappingType &first_kv = sib_page->ItemAt(0);
        page->InsertAt(page->GetSize(), first_kv);
        sib_page->Remove(0);
        parent->KeyAt(sib_page_id) = sib_page->KeyAt(0);
      } else {
        MappingType &last_kv = sib_page->Back();
        page->InsertAt(0, last_kv);
        sib_page->Remove(sib_page->GetSize());
        parent->KeyAt(page_id) = page->KeyAt(0);
      }
    }
  }
}

/*****************************************************************************
 * INDEX ITERATOR
 *****************************************************************************/
/*
 * Input parameter is void, find the leftmost leaf page first, then construct
 * index iterator
 * @return : index iterator
 */
INDEX_TEMPLATE_ARGUMENTS
auto BPLUSTREE_TYPE::Begin() -> INDEXITERATOR_TYPE {
  if (IsEmpty()) {
    return End();
  }
  page_id_t cur_id = GetRootPageId();
  while (true) {
    ReadPageGuard guard = bpm_->FetchPageRead(cur_id);
    auto base_page = guard.As<BPlusTreePage>();
    if (base_page->IsLeafPage()) {
      return INDEXITERATOR_TYPE(cur_id, 0, bpm_);
    } else {
      cur_id = guard.As<InternalPage>()->ValueAt(0);
    }
  }
}

/*
 * Input parameter is low key, find the leaf page that contains the input key
 * first, then construct index iterator
 * @return : index iterator
 */
INDEX_TEMPLATE_ARGUMENTS
auto BPLUSTREE_TYPE::Begin(const KeyType &key) -> INDEXITERATOR_TYPE {
  if (IsEmpty()) {
    return End();
  }
  std::vector<ValueType> result;
  Context ctx;
  if (FindLeafRead(ctx, key)) {
    auto base_page = ctx.read_set_.back().As<BPlusTreePage>();
    if (base_page->IsLeafPage()) {
      int idx = FindIdx(base_page, key);
      auto page = ctx.read_set_.back().As<LeafPage>();
      if (idx > 0 && comparator_(page->KeyAt(idx - 1), key) == 0) {
        while (ctx.read_set_.size() > 1) {
          ctx.read_set_.front().Drop();
          ctx.read_set_.pop_front();
        }
        return INDEXITERATOR_TYPE(ctx.read_set_.back().PageId(), idx - 1, bpm_);
      }
    }
  }
  CleanUp(ctx);
  return End();
}

/*
 * Input parameter is void, construct an index iterator representing the end
 * of the key/value pair in the leaf node
 * @return : index iterator
 */
INDEX_TEMPLATE_ARGUMENTS
auto BPLUSTREE_TYPE::End() -> INDEXITERATOR_TYPE { return INDEXITERATOR_TYPE(INVALID_PAGE_ID, -1, nullptr); }
/**
 * @return Page id of the root of this tree
 */
INDEX_TEMPLATE_ARGUMENTS
auto BPLUSTREE_TYPE::GetRootPageId() -> page_id_t {
  ReadPageGuard guard = bpm_->FetchPageRead(header_page_id_);
  auto header_page = guard.As<BPlusTreeHeaderPage>();
  return header_page->root_page_id_;
}

INDEX_TEMPLATE_ARGUMENTS
void BPLUSTREE_TYPE::InsertAtLeaf(Context &ctx, int index, const KeyType &key, const ValueType &value) {
  auto guard = std::move(ctx.write_set_.back());
  ctx.write_set_.pop_back();

  page_id_t orig_page_id = guard.PageId();
  auto page = guard.AsMut<LeafPage>();
  int min_size = page->GetMinSize();
  if (page->GetSize() + 1 < page->GetMaxSize()) {
    // early release
    ctx.header_page_ = std::nullopt;
    while (not ctx.write_set_.empty()) {
      ctx.write_set_.front().Drop();
      ctx.write_set_.pop_front();
    }
    page->InsertAt(index, MappingType(key, value));
    return;

  } else {
    page_id_t page_id;
    auto new_guard = bpm_->NewPageGuarded(&page_id);
    auto new_leaf_page = new_guard.AsMut<LeafPage>();
    new_leaf_page->Init(leaf_max_size_);
    new_leaf_page->SetNextPageId(page->GetNextPageId());

    page->SetNextPageId(page_id);
    page->InsertAt(index, MappingType(key, value));
    int moved_size = page->GetSize() - min_size;

    std::copy(page->Array() + min_size, page->End(), new_leaf_page->Begin());
    page->IncreaseSize(-moved_size);
    new_leaf_page->IncreaseSize(moved_size);

    if (ctx.write_set_.empty()) {
      NewRoot(ctx, orig_page_id, page_id, new_leaf_page->KeyAt(0));
    } else {
      InsertAtInternal(ctx, new_leaf_page->KeyAt(0), page_id);
    }
  }
}

INDEX_TEMPLATE_ARGUMENTS
void BPLUSTREE_TYPE::InsertAtInternal(Context &ctx, const KeyType &key, const page_id_t &value) {
  auto guard = std::move(ctx.write_set_.back());
  ctx.write_set_.pop_back();

  auto page = guard.AsMut<InternalPage>();
  page_id_t orig_page_id = guard.PageId();
  int index = FindIdx(page, key);

  if (page->GetSize() < page->GetMaxSize()) {
    // early release
    ctx.header_page_ = std::nullopt;
    while (not ctx.write_set_.empty()) {
      ctx.write_set_.front().Drop();
      ctx.write_set_.pop_front();
    }
    page->InsertAt(index, std::make_pair(key, value));
    return;
  } else {
    page_id_t page_id;
    auto new_guard = bpm_->NewPageGuarded(&page_id);
    auto new_page = new_guard.AsMut<InternalPage>();
    new_page->Init(internal_max_size_);

    int min_size = page->GetMinSize();
    if (index < min_size) {
      --min_size;
      int moved_size = page->GetSize() - min_size;
      std::copy(page->Array() + min_size, page->End(), new_page->Begin());
      page->IncreaseSize(-moved_size);
      new_page->IncreaseSize(moved_size);
      page->InsertAt(index, std::make_pair(key, value));
    } else {
      int moved_size = page->GetSize() - min_size;
      std::copy(page->Array() + min_size, page->End(), new_page->Begin());
      page->IncreaseSize(-moved_size);
      new_page->IncreaseSize(moved_size);
      new_page->InsertAt(index - min_size, std::make_pair(key, value));
    }

    if (ctx.write_set_.empty()) {
      NewRoot(ctx, orig_page_id, page_id, new_page->Array()[0].first);
    } else {
      InsertAtInternal(ctx, new_page->Array()[0].first, page_id);
    }
  }
}

INDEX_TEMPLATE_ARGUMENTS
void BPLUSTREE_TYPE::NewRoot(Context &ctx, const page_id_t &lpage_id, const page_id_t &rpage_id, const KeyType &key) {
  page_id_t new_page_id;
  auto new_guard = bpm_->NewPageGuarded(&new_page_id);
  auto new_page = new_guard.AsMut<InternalPage>();
  new_page->Init(internal_max_size_);
  new_page->InsertAt(0, std::make_pair(KeyType(), lpage_id));
  new_page->InsertAt(1, std::make_pair(key, rpage_id));
  ctx.header_page_->AsMut<BPlusTreeHeaderPage>()->root_page_id_ = new_page_id;
}

INDEX_TEMPLATE_ARGUMENTS
auto BPLUSTREE_TYPE::FindIdx(const BPlusTreePage *base_page, const KeyType &key) -> int {
  if (base_page->IsLeafPage()) {
    auto page = reinterpret_cast<const LeafPage *>(base_page);
    int count = base_page->GetSize();
    int l = 0;
    while (count > 0) {  // upper_bound
      int step = count / 2;
      int it = l + step;
      if (comparator_(page->KeyAt(it), key) <= 0) {
        l = ++it;
        count -= step + 1;
      } else {
        count = step;
      }
    }
    return l;
  }
  auto page = reinterpret_cast<const InternalPage *>(base_page);
  int l = 1, count = page->GetSize() - 1;
  while (count > 0) {  // upper_bound
    int step = count / 2;
    int it = l + step;
    if (comparator_(page->KeyAt(it), key) <= 0) {
      l = ++it;
      count -= step + 1;
    } else {
      count = step;
    }
  }
  return l;
}

INDEX_TEMPLATE_ARGUMENTS
auto BPLUSTREE_TYPE::FindLeafWrite(Context &ctx, const KeyType &key, bool is_insert) -> bool {
  WritePageGuard header_guard = bpm_->FetchPageWrite(header_page_id_);
  ctx.root_page_id_ = header_guard.As<BPlusTreeHeaderPage>()->root_page_id_;
  ctx.header_page_ = std::move(header_guard);
  page_id_t cur_page_id = ctx.root_page_id_;

  while (true) {
    if (cur_page_id == INVALID_PAGE_ID) {
      return false;
    }
    WritePageGuard guard = bpm_->FetchPageWrite(cur_page_id);
    auto base_page = guard.As<BPlusTreePage>();

    if (base_page->IsLeafPage()) {
      ctx.write_set_.emplace_back(std::move(guard));
      return true;
    } else {
      int idx = FindIdx(base_page, key);
      auto page = guard.As<InternalPage>();
      if ((not ctx.write_set_.empty()) && ((is_insert && page->GetSize() + 1 < page->GetMaxSize()) or
                                           (!is_insert && page->GetSize() > page->GetMinSize()))) {
        CleanUp(ctx);
      }
      cur_page_id = page->ValueAt(idx - 1);
      ctx.write_set_.emplace_back(std::move(guard));
    }
  }
}
INDEX_TEMPLATE_ARGUMENTS
auto BPLUSTREE_TYPE::OptimisticFindLeaf(Context &ctx, const KeyType &key) -> bool {
  ctx.root_page_id_ = GetRootPageId();

  page_id_t cur_page_id = ctx.root_page_id_;

  while (true) {
    if (cur_page_id == INVALID_PAGE_ID) {
      return false;
    }
    ReadPageGuard guard = bpm_->FetchPageRead(cur_page_id);
    auto base_page = guard.As<BPlusTreePage>();

    if (base_page->IsLeafPage()) {
      guard.Drop();
      ctx.write_set_.emplace_back(bpm_->FetchPageWrite(cur_page_id));
      return true;
    } else {
      int idx = FindIdx(base_page, key);
      auto page = guard.As<InternalPage>();
      cur_page_id = page->ValueAt(idx - 1);
      ctx.read_set_.emplace_back(std::move(guard));
    }
    while (ctx.read_set_.size() > 1) {
      ctx.read_set_.front().Drop();
      ctx.read_set_.pop_front();
    }
  }
}

INDEX_TEMPLATE_ARGUMENTS
auto BPLUSTREE_TYPE::FindLeafRead(Context &ctx, const KeyType &key) -> bool {
  ctx.root_page_id_ = GetRootPageId();
  page_id_t cur_page_id = ctx.root_page_id_;

  while (true) {
    if (cur_page_id == INVALID_PAGE_ID) {
      return false;
    }
    ReadPageGuard guard = bpm_->FetchPageRead(cur_page_id);
    auto base_page = guard.As<BPlusTreePage>();

    if (base_page->IsLeafPage()) {
      ctx.read_set_.emplace_back(std::move(guard));
      return true;
    } else {
      int idx = FindIdx(base_page, key);
      auto page = guard.As<InternalPage>();
      cur_page_id = page->ValueAt(idx - 1);
      ctx.read_set_.emplace_back(std::move(guard));
    }
    while (ctx.read_set_.size() > 1) {
      ctx.read_set_.front().Drop();
      ctx.read_set_.pop_front();
    }
  }
}
/*****************************************************************************
 * UTILITIES AND DEBUG
 *****************************************************************************/

/*
 * This method is used for test only
 * Read data from file and insert one by one
 */
INDEX_TEMPLATE_ARGUMENTS
void BPLUSTREE_TYPE::InsertFromFile(const std::string &file_name, Transaction *txn) {
  int64_t key;
  std::ifstream input(file_name);
  while (input) {
    input >> key;

    KeyType index_key;
    index_key.SetFromInteger(key);
    RID rid(key);
    Insert(index_key, rid, txn);
  }
}
/*
 * This method is used for test only
 * Read data from file and remove one by one
 */
INDEX_TEMPLATE_ARGUMENTS
void BPLUSTREE_TYPE::RemoveFromFile(const std::string &file_name, Transaction *txn) {
  int64_t key;
  std::ifstream input(file_name);
  while (input) {
    input >> key;
    KeyType index_key;
    index_key.SetFromInteger(key);
    Remove(index_key, txn);
  }
}

INDEX_TEMPLATE_ARGUMENTS
void BPLUSTREE_TYPE::Print(BufferPoolManager *bpm) {
  auto root_page_id = GetRootPageId();
  auto guard = bpm->FetchPageBasic(root_page_id);
  PrintTree(guard.PageId(), guard.template As<BPlusTreePage>());
}

INDEX_TEMPLATE_ARGUMENTS
void BPLUSTREE_TYPE::PrintTree(page_id_t page_id, const BPlusTreePage *page) {
  if (page->IsLeafPage()) {
    auto *leaf = reinterpret_cast<const LeafPage *>(page);
    std::cout << "Leaf Page: " << page_id << "\tNext: " << leaf->GetNextPageId() << std::endl;

    // Print the contents of the leaf page.
    std::cout << "Contents: ";
    for (int i = 0; i < leaf->GetSize(); i++) {
      std::cout << leaf->KeyAt(i);
      if ((i + 1) < leaf->GetSize()) {
        std::cout << ", ";
      }
    }
    std::cout << std::endl;
    std::cout << std::endl;

  } else {
    auto *internal = reinterpret_cast<const InternalPage *>(page);
    std::cout << "Internal Page: " << page_id << std::endl;

    // Print the contents of the internal page.
    std::cout << "Contents: ";
    for (int i = 0; i < internal->GetSize(); i++) {
      std::cout << internal->KeyAt(i) << ": " << internal->ValueAt(i);
      if ((i + 1) < internal->GetSize()) {
        std::cout << ", ";
      }
    }
    std::cout << std::endl;
    std::cout << std::endl;
    for (int i = 0; i < internal->GetSize(); i++) {
      auto guard = bpm_->FetchPageBasic(internal->ValueAt(i));
      PrintTree(guard.PageId(), guard.template As<BPlusTreePage>());
    }
  }
}

/**
 * This method is used for debug only, You don't need to modify
 */
INDEX_TEMPLATE_ARGUMENTS
void BPLUSTREE_TYPE::Draw(BufferPoolManager *bpm, const std::string &outf) {
  if (IsEmpty()) {
    LOG_WARN("Drawing an empty tree");
    return;
  }

  std::ofstream out(outf);
  out << "digraph G {" << std::endl;
  auto root_page_id = GetRootPageId();
  auto guard = bpm->FetchPageBasic(root_page_id);
  ToGraph(guard.PageId(), guard.template As<BPlusTreePage>(), out);
  out << "}" << std::endl;
  out.close();
}

/**
 * This method is used for debug only, You don't need to modify
 */
INDEX_TEMPLATE_ARGUMENTS
void BPLUSTREE_TYPE::ToGraph(page_id_t page_id, const BPlusTreePage *page, std::ofstream &out) {
  std::string leaf_prefix("LEAF_");
  std::string internal_prefix("INT_");
  if (page->IsLeafPage()) {
    auto *leaf = reinterpret_cast<const LeafPage *>(page);
    // Print node name
    out << leaf_prefix << page_id;
    // Print node properties
    out << "[shape=plain color=green ";
    // Print data of the node
    out << "label=<<TABLE BORDER=\"0\" CELLBORDER=\"1\" CELLSPACING=\"0\" CELLPADDING=\"4\">\n";
    // Print data
    out << "<TR><TD COLSPAN=\"" << leaf->GetSize() << "\">P=" << page_id << "</TD></TR>\n";
    out << "<TR><TD COLSPAN=\"" << leaf->GetSize() << "\">"
        << "max_size=" << leaf->GetMaxSize() << ",min_size=" << leaf->GetMinSize() << ",size=" << leaf->GetSize()
        << "</TD></TR>\n";
    out << "<TR>";
    for (int i = 0; i < leaf->GetSize(); i++) {
      out << "<TD>" << leaf->KeyAt(i) << "</TD>\n";
    }
    out << "</TR>";
    // Print table end
    out << "</TABLE>>];\n";
    // Print Leaf node link if there is a next page
    if (leaf->GetNextPageId() != INVALID_PAGE_ID) {
      out << leaf_prefix << page_id << " -> " << leaf_prefix << leaf->GetNextPageId() << ";\n";
      out << "{rank=same " << leaf_prefix << page_id << " " << leaf_prefix << leaf->GetNextPageId() << "};\n";
    }
  } else {
    auto *inner = reinterpret_cast<const InternalPage *>(page);
    // Print node name
    out << internal_prefix << page_id;
    // Print node properties
    out << "[shape=plain color=pink ";  // why not?
    // Print data of the node
    out << "label=<<TABLE BORDER=\"0\" CELLBORDER=\"1\" CELLSPACING=\"0\" CELLPADDING=\"4\">\n";
    // Print data
    out << "<TR><TD COLSPAN=\"" << inner->GetSize() << "\">P=" << page_id << "</TD></TR>\n";
    out << "<TR><TD COLSPAN=\"" << inner->GetSize() << "\">"
        << "max_size=" << inner->GetMaxSize() << ",min_size=" << inner->GetMinSize() << ",size=" << inner->GetSize()
        << "</TD></TR>\n";
    out << "<TR>";
    for (int i = 0; i < inner->GetSize(); i++) {
      out << "<TD PORT=\"p" << inner->ValueAt(i) << "\">";
      if (i > 0) {
        out << inner->KeyAt(i);
      } else {
        out << " ";
      }
      out << "</TD>\n";
    }
    out << "</TR>";
    // Print table end
    out << "</TABLE>>];\n";
    // Print leaves
    for (int i = 0; i < inner->GetSize(); i++) {
      auto child_guard = bpm_->FetchPageBasic(inner->ValueAt(i));
      auto child_page = child_guard.template As<BPlusTreePage>();
      ToGraph(child_guard.PageId(), child_page, out);
      if (i > 0) {
        auto sibling_guard = bpm_->FetchPageBasic(inner->ValueAt(i - 1));
        auto sibling_page = sibling_guard.template As<BPlusTreePage>();
        if (!sibling_page->IsLeafPage() && !child_page->IsLeafPage()) {
          out << "{rank=same " << internal_prefix << sibling_guard.PageId() << " " << internal_prefix
              << child_guard.PageId() << "};\n";
        }
      }
      out << internal_prefix << page_id << ":p" << child_guard.PageId() << " -> ";
      if (child_page->IsLeafPage()) {
        out << leaf_prefix << child_guard.PageId() << ";\n";
      } else {
        out << internal_prefix << child_guard.PageId() << ";\n";
      }
    }
  }
}

INDEX_TEMPLATE_ARGUMENTS
auto BPLUSTREE_TYPE::DrawBPlusTree() -> std::string {
  if (IsEmpty()) {
    return "()";
  }

  PrintableBPlusTree p_root = ToPrintableBPlusTree(GetRootPageId());
  std::ostringstream out_buf;
  p_root.Print(out_buf);

  return out_buf.str();
}

INDEX_TEMPLATE_ARGUMENTS
auto BPLUSTREE_TYPE::ToPrintableBPlusTree(page_id_t root_id) -> PrintableBPlusTree {
  auto root_page_guard = bpm_->FetchPageBasic(root_id);
  auto root_page = root_page_guard.template As<BPlusTreePage>();
  PrintableBPlusTree proot;

  if (root_page->IsLeafPage()) {
    auto leaf_page = root_page_guard.template As<LeafPage>();
    proot.keys_ = leaf_page->ToString();
    proot.size_ = proot.keys_.size() + 4;  // 4 more spaces for indent

    return proot;
  }

  // draw internal page
  auto internal_page = root_page_guard.template As<InternalPage>();
  proot.keys_ = internal_page->ToString();
  proot.size_ = 0;
  for (int i = 0; i < internal_page->GetSize(); i++) {
    page_id_t child_id = internal_page->ValueAt(i);
    PrintableBPlusTree child_node = ToPrintableBPlusTree(child_id);
    proot.size_ += child_node.size_;
    proot.children_.push_back(child_node);
  }

  return proot;
}

template class BPlusTree<GenericKey<4>, RID, GenericComparator<4>>;

template class BPlusTree<GenericKey<8>, RID, GenericComparator<8>>;

template class BPlusTree<GenericKey<16>, RID, GenericComparator<16>>;

template class BPlusTree<GenericKey<32>, RID, GenericComparator<32>>;

template class BPlusTree<GenericKey<64>, RID, GenericComparator<64>>;

}  // namespace bustub
